import torch, os, cv2
from model.model import parsingNet
from utils1.common import merge_config
from utils1.dist_utils import dist_print
import torch
import scipy.special, tqdm
import numpy as np
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset
from data.constant import culane_row_anchor, tusimple_row_anchor
import sys
import time
from PIL import Image
project_path = 'E:/pythonprojectnew/graduation/Ultra-Fast-Lane-Detection-master'
sys.path.append(project_path)
img_path=r'E:\pythonprojectnew\graduation\Ultra-Fast-Lane-Detection-master\CUlane\driver_100_30frame\05251548_0439.MP4\00450.JPG'
if __name__ == "__main__":

    cv2.namedWindow("vis", 0)
    cv2.resizeWindow("vis", 640, 480)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.backends.cudnn.benchmark = True

    args, cfg = merge_config()

    dist_print('start testing...')
    assert cfg.backbone in ['18', '34', '50', '101', '152', '50next', '101next', '50wide', '101wide']

    if cfg.dataset == 'CULane':
        cls_num_per_lane = 18
    elif cfg.dataset == 'Tusimple':
        cls_num_per_lane = 56
    else:
        raise NotImplementedError

    # net = parsingNet(pretrained = False, backbone=cfg.backbone,cls_dim = (cfg.griding_num+1,cls_num_per_lane,4),
    #                 use_aux=False).cuda() # we dont need auxiliary segmentation in testing

    net = parsingNet(pretrained=False, backbone=cfg.backbone, cls_dim=(cfg.griding_num + 1, cls_num_per_lane, 4),
                     use_aux=False).to(device)  # we dont need auxiliary segmentation in testing

    state_dict = torch.load(cfg.test_model, map_location='cpu')['model']
    compatible_state_dict = {}
    for k, v in state_dict.items():
        if 'module.' in k:
            compatible_state_dict[k[7:]] = v
        else:
            compatible_state_dict[k] = v

    net.load_state_dict(compatible_state_dict, strict=False)
    net.eval()

    img_transforms = transforms.Compose([
        transforms.Resize((288, 800)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    if cfg.dataset == 'CULane':
        splits = ['test0_normal.txt', 'test1_crowd.txt', 'test2_hlight.txt', 'test3_shadow.txt', 'test4_noline.txt',
                  'test5_arrow.txt', 'test6_curve.txt', 'test7_cross.txt', 'test8_night.txt']
        datasets = [LaneTestDataset(cfg.data_root, os.path.join(cfg.data_root, 'list/test_split/' + split),
                                    img_transform=img_transforms) for split in splits]
        img_w, img_h = 1640, 590
        row_anchor = culane_row_anchor
    elif cfg.dataset == 'Tusimple':
        splits = ['test.txt']
        datasets = [LaneTestDataset(cfg.data_root, os.path.join(cfg.data_root, split), img_transform=img_transforms) for
                    split in splits]
        img_w, img_h = 1280, 720
        row_anchor = tusimple_row_anchor
    else:
        raise NotImplementedError
    for split, dataset in zip(splits, datasets):
        print("split:", split)
        print("dataset:", type(dataset))

        img_cv = Image.open(img_path)
        img_cv = img_transforms(img_cv)
        img_cv=img_cv.unsqueeze(0)
        print("?",type(img_cv),img_cv.shape)
        #img_cv = torch.from_numpy(img_cv)


        loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
        fourcc = cv2.VideoWriter_fourcc(*'MJPG')
        # print("这是什么",split[:-3]+'avi')
        # print("loader:",type(loader),loader)
        # print("img_cv:",type(img_cv),img_cv)
        vout = cv2.VideoWriter(split[:-3] + 'avi', fourcc, 30.0, (img_w, img_h))
        for i, data in enumerate(tqdm.tqdm(loader)):
            print(type(data))
            print(data)
            imgs, names = data
            # t_start = time.time()
            print("imgs:",type(imgs),imgs.shape)
            print("img_cv:",type(img_cv),img_cv.shape)
            # imgs = imgs.cuda()
            imgs = imgs.to(device)

            with torch.no_grad():
                out = net(imgs)
            print("out:",type(out),out.shape)
            # print("运行时间：", time.time() - t_start)
            col_sample = np.linspace(0, 800 - 1, cfg.griding_num)
            col_sample_w = col_sample[1] - col_sample[0]

            out_j = out[0].data.cpu().numpy()
            out_j = out_j[:, ::-1, :]
            prob = scipy.special.softmax(out_j[:-1, :, :], axis=0)
            idx = np.arange(cfg.griding_num) + 1
            idx = idx.reshape(-1, 1, 1)
            loc = np.sum(prob * idx, axis=0)
            out_j = np.argmax(out_j, axis=0)
            loc[out_j == cfg.griding_num] = 0
            out_j = loc

            # import pdb; pdb.set_trace()
            vis = cv2.imread(os.path.join(cfg.data_root, names[0]))
            vis = cv2.resize(vis, (img_w, img_h), cv2.INTER_AREA)
            for i in range(out_j.shape[1]):
                print(i)
                if np.sum(out_j[:, i] != 0) > 2:
                    for k in range(out_j.shape[0]):
                        if out_j[k, i] > 0:
                            ppp = (int(out_j[k, i] * col_sample_w * img_w / 800) - 1,
                                   int(img_h * (row_anchor[cls_num_per_lane - 1 - k] / 288)) - 1)
                            cv2.circle(vis, ppp, 5, (0, 255, 0), -1)
                            # print("ppp",ppp)
                cv2.imshow('vis', vis)
                cv2.waitKey(0)
                cv2.destroyAllWindows()

            vout.write(vis)

        vout.release()